#include "kernel_operator.h"

using namespace AscendC;

constexpr int32_t BUFFER_NUM = 1;

template <typename T>
class TokenShift
{
public:
    __aicore__ inline TokenShift() {}
    __aicore__ inline void Init(
        GM_ADDR x, 
        GM_ADDR rwkvag, 
        GM_ADDR h0, 
        GM_ADDR xr, 
        GM_ADDR xw, 
        GM_ADDR xk,
        GM_ADDR xv, 
        GM_ADDR xa, 
        GM_ADDR xg, 
        GM_ADDR ht, 
        GM_ADDR workspace,
        uint32_t B_length,
        uint32_t T_length,
        uint32_t C_length,
        uint32_t tiling_size,
        uint32_t core_rows,
        uint32_t core_rows_remain,
        uint32_t X_length_align,
        TPipe* pipeIn)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero");
        this->pipe = pipeIn;
        this->startPointer = core_rows * T_length * C_length * GetBlockIdx();
        this->core_rows = core_rows + (GetBlockNum() == GetBlockIdx() + 1 ? core_rows_remain : 0);
        this->B_length = B_length;
        this->T_length = T_length;
        this->C_length = C_length;
        this->tiling_size = tiling_size;
        this->ALIGN_NUM = 32 / sizeof(T);
        uint32_t core_length = this->core_rows * T_length * C_length + ALIGN_NUM;
        uint32_t h0startPointer = core_rows * C_length * GetBlockIdx();

        x_global.SetGlobalBuffer((__gm__ T *)x + this->startPointer, core_length);
        rwkvag_global.SetGlobalBuffer((__gm__ T *)rwkvag, 6 * C_length + ALIGN_NUM);
        h0_global.SetGlobalBuffer((__gm__ T *)h0 + h0startPointer, this->core_rows * C_length + ALIGN_NUM);

        xr_global.SetGlobalBuffer((__gm__ T *)xr + this->startPointer, core_length);   
        xw_global.SetGlobalBuffer((__gm__ T *)xw + this->startPointer, core_length);   
        xk_global.SetGlobalBuffer((__gm__ T *)xk + this->startPointer, core_length);   
        xv_global.SetGlobalBuffer((__gm__ T *)xv + this->startPointer, core_length);   
        xa_global.SetGlobalBuffer((__gm__ T *)xa + this->startPointer, core_length);   
        xg_global.SetGlobalBuffer((__gm__ T *)xg + this->startPointer, core_length);        
        ht_global.SetGlobalBuffer((__gm__ T *)ht + h0startPointer, B_length * C_length + ALIGN_NUM); 
        pipe->InitBuffer(inQueueXBind, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(inQueueX, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(inQueueRwkvag, BUFFER_NUM, 6 * this->tiling_size * sizeof(T));
        pipe->InitBuffer(inQueueH0, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(inQueueXNextLine, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXR, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXW, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXK, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXV, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXA, BUFFER_NUM, this->tiling_size * sizeof(T));
        pipe->InitBuffer(outQueueXG, BUFFER_NUM, this->tiling_size * sizeof(T));

        pipe->InitBuffer(QueueTmp1, this->tiling_size * sizeof(T));
        g_xLocal = QueueTmp1.Get<T>();
    }

    __aicore__ inline void Process() {
        int32_t loopcount = this->C_length / this->tiling_size;
        uint32_t length_remain = this->C_length % this->tiling_size;
        uint32_t length_remain_align = (length_remain + this->ALIGN_NUM - 1) / this->ALIGN_NUM * this->ALIGN_NUM;

        for(int32_t p = 0; p < loopcount; p++) {
            for(int32_t k = 0; k < this->core_rows; k++) {
                CopyInH0(k * this->C_length + p * this->tiling_size, this->tiling_size);
                CopyInX(k * this->T_length * this->C_length + p * this->tiling_size, this->tiling_size);
                ComputeTopLine(p * this->tiling_size, this->tiling_size);
                CopyOutX(k * this->T_length * this->C_length + p * this->tiling_size, this->tiling_size);

                for(int32_t m = 1; m < this->T_length; m++) {
                    
                    CopyInXNextLine(k * this->T_length * this->C_length + m * this->C_length + p * this->tiling_size, this->tiling_size);
                    Compute(p * this->tiling_size, this->tiling_size);
                    CopyOutX(k * this->T_length * this->C_length + m * this->C_length + p * this->tiling_size, this->tiling_size);
                    
                }  
                CopyOutHT(k * this->C_length + p * this->tiling_size, this->tiling_size);
            }    
        }

        if(length_remain > 0) {
            for(int32_t k = 0; k < this->core_rows; k++) {
                CopyInH0(k * this->C_length + loopcount * this->tiling_size, length_remain_align);
                CopyInX(k * this->T_length * this->C_length + loopcount * this->tiling_size, length_remain_align);
                ComputeTopLine(loopcount * this->tiling_size, length_remain_align);
                // if(length_remain == length_remain_align) {
                //     CopyOutX(k * this->T_length * this->C_length + loopcount * this->tiling_size, length_remain_align);
                // } else {
                //     CopyOutXUnalign(k * this->T_length * this->C_length + loopcount * this->tiling_size, length_remain);
                // }
                CopyOutXTail(k * this->T_length * this->C_length + loopcount * this->tiling_size, length_remain);
                for(int32_t m = 1; m < this->T_length; m++) {
                    CopyInXNextLine(k * this->T_length * this->C_length + m * this->C_length + loopcount * this->tiling_size, length_remain_align);
                    Compute(loopcount * this->tiling_size, length_remain_align);
                    // if(length_remain == length_remain_align) {
                    //     CopyOutX(k * this->T_length * this->C_length + m * this->C_length + loopcount * this->tiling_size, length_remain_align);
                    // } else {
                    //     CopyOutXUnalign(k * this->T_length * this->C_length + m * this->C_length + loopcount * this->tiling_size, length_remain);
                    // }
                    CopyOutXTail(k * this->T_length * this->C_length + m * this->C_length + loopcount * this->tiling_size, length_remain);
                }
                CopyOutHTTail(k * this->C_length + loopcount * this->tiling_size, length_remain);
            }    
        }
               
    }
    __aicore__ inline void ProcessHT(GM_ADDR x)
    {
        if(GetBlockIdx() == 0) {
            // ht 
            x_global.SetGlobalBuffer((__gm__ T *)x, this->X_length_align);
            if(this->T_length == 1) {
                int32_t loopcount = this->B_length * this->C_length / this->tiling_size;
                for(int32_t k = 0; k < loopcount; k++) {
                    CopyInHT(k * this->tiling_size, this->tiling_size);
                    CopyOutHT(k * this->tiling_size, this->tiling_size);
                }
                uint32_t length_remain = this->B_length * this->C_length % this->tiling_size;
                if(length_remain > 0) {
                    uint32_t process_length = length_remain / this->ALIGN_NUM * this->ALIGN_NUM;
                    CopyInHT(loopcount * this->tiling_size, process_length);
                    CopyOutHT(loopcount * this->tiling_size, process_length);
                    for(int32_t m = loopcount * this->tiling_size + process_length; m < this->B_length * this->C_length; m++) {
                        ht_global.SetValue(m, x_global.GetValue(m));
                        PipeBarrier<PIPE_ALL>();
                    }
                }                 
            } else {
                for(int32_t k = 0; k < this->B_length; k++) {
                    int32_t loopcount = this->C_length / this->tiling_size;
                    for(int32_t m = 0; m < loopcount; m++) {
                        CopyInHT(k * this->T_length * this->C_length + (this->T_length - 1) * this->C_length + m * this->tiling_size, this->tiling_size);
                        CopyOutHT(k * this->C_length + m * this->tiling_size, this->tiling_size);
                    }
                    uint32_t length_remain = this->C_length % this->tiling_size;
                    if(length_remain > 0) {
                        uint32_t process_length = length_remain / this->ALIGN_NUM * this->ALIGN_NUM;
                        CopyInHT(k * this->T_length * this->C_length + (this->T_length - 1) * this->C_length + loopcount * this->tiling_size, process_length);
                        CopyOutHT(k * this->C_length + loopcount * this->tiling_size, process_length);
                        for(int32_t m = loopcount * this->tiling_size + process_length; m < this->C_length; m++) {
                            ht_global.SetValue(m + (k * this->C_length), x_global.GetValue(m + (k * this->T_length * this->C_length) + ((this->T_length - 1) * this->C_length)));
                            PipeBarrier<PIPE_ALL>();
                        }
                    } 
                }
            }            
        }

        // int32_t loopCount = this->tileNum;
        // for (int32_t i = 0; i < loopCount - 1; i++)
        // {
        //     CopyIn(i, this->tileLength);
        //     Compute(i, this->tileLength);
        //     CopyOut(i, this->tileLength);
        // }
    }    
    __aicore__ inline void ComputeTopLine(uint32_t startpos, uint32_t length) {
        LocalTensor<T> h0Local = inQueueH0.DeQue<T>();
        // LocalTensor<T> xLocal = inQueueX.DeQue<T>();
        LocalTensor<T> XRLocal = outQueueXR.AllocTensor<T>();
        LocalTensor<T> XWLocal = outQueueXW.AllocTensor<T>();
        LocalTensor<T> XKLocal = outQueueXK.AllocTensor<T>();
        LocalTensor<T> XVLocal = outQueueXV.AllocTensor<T>();
        LocalTensor<T> XALocal = outQueueXA.AllocTensor<T>();
        LocalTensor<T> XGLocal = outQueueXG.AllocTensor<T>();

        Sub(h0Local, h0Local, g_xLocal, length);

        CopyInRline(startpos, length);
        LocalTensor<T> rwkvagLocal = inQueueRwkvag.DeQue<T>();
        Mul(XRLocal, h0Local, rwkvagLocal, length);
        Add(XRLocal, XRLocal, g_xLocal, length);
        // inQueueRwkvag.FreeTensor(rwkvagLocal0);
        outQueueXR.EnQue(XRLocal);  

        // CopyInRline(startpos + this->C_length, length);
        // LocalTensor<T> rwkvagLocal1 = inQueueRwkvag.DeQue<T>();
        Mul(XWLocal, h0Local, rwkvagLocal[this->tiling_size], length);
        Add(XWLocal, XWLocal, g_xLocal, length);
        // inQueueRwkvag.FreeTensor(rwkvagLocal1);
        outQueueXW.EnQue(XWLocal);  
        
        // CopyInRline(startpos + this->C_length * 2, length);
        // LocalTensor<T> rwkvagLocal2 = inQueueRwkvag.DeQue<T>();
        Mul(XKLocal, h0Local, rwkvagLocal[this->tiling_size * 2], length);
        Add(XKLocal, XKLocal, g_xLocal, length);
        // inQueueRwkvag.FreeTensor(rwkvagLocal2);
        outQueueXK.EnQue(XKLocal);  

        // CopyInRline(startpos + this->C_length * 3, length);
        // LocalTensor<T> rwkvagLocal3 = inQueueRwkvag.DeQue<T>();
        Mul(XVLocal, h0Local, rwkvagLocal[this->tiling_size * 3], length);
        Add(XVLocal, XVLocal, g_xLocal, length);
        // inQueueRwkvag.FreeTensor(rwkvagLocal3);
        outQueueXV.EnQue(XVLocal); 

        // CopyInRline(startpos + this->C_length * 4, length);
        // LocalTensor<T> rwkvagLocal4 = inQueueRwkvag.DeQue<T>();
        Mul(XALocal, h0Local, rwkvagLocal[this->tiling_size * 4], length);
        Add(XALocal, XALocal, g_xLocal, length);
        // inQueueRwkvag.FreeTensor(rwkvagLocal4);
        outQueueXA.EnQue(XALocal); 

        // CopyInRline(startpos + this->C_length * 5, length);
        // LocalTensor<T> rwkvagLocal5 = inQueueRwkvag.DeQue<T>();
        Mul(XGLocal, h0Local, rwkvagLocal[this->tiling_size * 5], length);
        Add(XGLocal, XGLocal, g_xLocal, length);
        inQueueRwkvag.FreeTensor(rwkvagLocal);
        outQueueXG.EnQue(XGLocal);      
        
        inQueueH0.FreeTensor(h0Local);
        // inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void Compute(uint32_t startpos, uint32_t length) {
        LocalTensor<T> xNextLineLocal = inQueueXNextLine.DeQue<T>();
        // LocalTensor<T> xLocal = inQueueX.DeQue<T>();
        LocalTensor<T> XRLocal = outQueueXR.AllocTensor<T>();
        LocalTensor<T> XWLocal = outQueueXW.AllocTensor<T>();
        LocalTensor<T> XKLocal = outQueueXK.AllocTensor<T>();
        LocalTensor<T> XVLocal = outQueueXV.AllocTensor<T>();
        LocalTensor<T> XALocal = outQueueXA.AllocTensor<T>();
        LocalTensor<T> XGLocal = outQueueXG.AllocTensor<T>();

        Sub(g_xLocal, g_xLocal, xNextLineLocal, length);

        CopyInRline(startpos, length);
        LocalTensor<T> rwkvagLocal = inQueueRwkvag.DeQue<T>();
        Mul(XRLocal, g_xLocal, rwkvagLocal, length);
        Add(XRLocal, XRLocal, xNextLineLocal, length);
        // inQueueRwkvag.FreeTensor(rwkvagLocal0);
        outQueueXR.EnQue(XRLocal);  

        // CopyInRline(startpos + this->C_length, length);
        // LocalTensor<T> rwkvagLocal1 = inQueueRwkvag.DeQue<T>();
        Mul(XWLocal, g_xLocal, rwkvagLocal[this->tiling_size], length);
        Add(XWLocal, XWLocal, xNextLineLocal, length);
        // inQueueRwkvag.FreeTensor(rwkvagLocal1);
        outQueueXW.EnQue(XWLocal);  
        
        // CopyInRline(startpos + this->C_length * 2, length);
        // LocalTensor<T> rwkvagLocal2 = inQueueRwkvag.DeQue<T>();
        Mul(XKLocal, g_xLocal, rwkvagLocal[this->tiling_size * 2], length);
        Add(XKLocal, XKLocal, xNextLineLocal, length);
        // inQueueRwkvag.FreeTensor(rwkvagLocal2);
        outQueueXK.EnQue(XKLocal);  

        // CopyInRline(startpos + this->C_length * 3, length);
        // LocalTensor<T> rwkvagLocal3 = inQueueRwkvag.DeQue<T>();
        Mul(XVLocal, g_xLocal, rwkvagLocal[this->tiling_size * 3], length);
        Add(XVLocal, XVLocal, xNextLineLocal, length);
        // inQueueRwkvag.FreeTensor(rwkvagLocal3);
        outQueueXV.EnQue(XVLocal); 

        // CopyInRline(startpos + this->C_length * 4, length);
        // LocalTensor<T> rwkvagLocal4 = inQueueRwkvag.DeQue<T>();
        Mul(XALocal, g_xLocal, rwkvagLocal[this->tiling_size * 4], length);
        Add(XALocal, XALocal, xNextLineLocal, length);
        // inQueueRwkvag.FreeTensor(rwkvagLocal4);
        outQueueXA.EnQue(XALocal); 

        // CopyInRline(startpos + this->C_length * 5, length);
        // LocalTensor<T> rwkvagLocal5 = inQueueRwkvag.DeQue<T>();
        Mul(XGLocal, g_xLocal, rwkvagLocal[this->tiling_size * 5], length);
        Add(XGLocal, XGLocal, xNextLineLocal, length);
        inQueueRwkvag.FreeTensor(rwkvagLocal);
        outQueueXG.EnQue(XGLocal);      
        
        Adds(g_xLocal, xNextLineLocal, T(0), length);

        // inQueueX.EnQue(xLocal);
        inQueueXNextLine.FreeTensor(xNextLineLocal);
        PipeBarrier<PIPE_ALL>(); 
    }    

    __aicore__ inline void CopyInH0(uint32_t startpos, uint32_t length) {
        LocalTensor<T> h0Local = inQueueH0.AllocTensor<T>();
        DataCopy(h0Local, h0_global[startpos], length);   
        inQueueH0.EnQue(h0Local);         
    }       
    __aicore__ inline void CopyInX(uint32_t startpos, uint32_t length) {
        // LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
        DataCopy(g_xLocal, x_global[startpos], length);  
        PipeBarrier<PIPE_ALL>(); 
        // inQueueX.EnQue(xLocal);  
    }

    __aicore__ inline void CopyInXNextLine(uint32_t startpos, uint32_t length) {
        LocalTensor<T> xLocal = inQueueXNextLine.AllocTensor<T>();
        DataCopy(xLocal, x_global[startpos], length);   
        inQueueXNextLine.EnQue(xLocal); 
        PipeBarrier<PIPE_ALL>();  
    }

    __aicore__ inline void CopyInRline(uint32_t startpos, uint32_t length) {
        LocalTensor<T> rwkvagLocal = inQueueRwkvag.AllocTensor<T>();
        DataCopy(rwkvagLocal, rwkvag_global[startpos], length);
        DataCopy(rwkvagLocal[this->tiling_size], rwkvag_global[startpos + this->C_length], length);  
        DataCopy(rwkvagLocal[this->tiling_size * 2], rwkvag_global[startpos + this->C_length * 2], length); 
        DataCopy(rwkvagLocal[this->tiling_size * 3], rwkvag_global[startpos + this->C_length * 3], length); 
        DataCopy(rwkvagLocal[this->tiling_size * 4], rwkvag_global[startpos + this->C_length * 4], length); 
        DataCopy(rwkvagLocal[this->tiling_size * 5], rwkvag_global[startpos + this->C_length * 5], length); 
        inQueueRwkvag.EnQue(rwkvagLocal); 
        PipeBarrier<PIPE_ALL>();  
    }
    __aicore__ inline void CopyInHT(uint32_t startpos, uint32_t length) {
        LocalTensor<T> xLocal = inQueueXBind.AllocTensor<T>();
        DataCopy(xLocal, x_global[startpos], length);   
        inQueueXBind.EnQue(xLocal);
        PipeBarrier<PIPE_ALL>();
    }

    // __aicore__ inline void CopyOutHT(uint32_t startpos, uint32_t length) {
    //     LocalTensor<T> xLocal = inQueueXBind.DeQue<T>();
    //     DataCopy(ht_global[startpos], xLocal, length);  
    //     inQueueXBind.FreeTensor(xLocal);
    //     PipeBarrier<PIPE_ALL>();
    // }
    __aicore__ inline void CopyOutHT(uint32_t startpos, uint32_t length) {
        // LocalTensor<T> xLocal = inQueueXBind.DeQue<T>();
        DataCopy(ht_global[startpos], g_xLocal, length);  
        PipeBarrier<PIPE_ALL>();  
    }
    __aicore__ inline void CopyOutHTTail(uint32_t startpos, uint32_t length) {
        for(int32_t k = 0; k < length; k++) {
            ht_global.SetValue(k + startpos, g_xLocal.GetValue(k));
        }
        PipeBarrier<PIPE_ALL>(); 
    }
    __aicore__ inline void CopyOutX(uint32_t startpos, uint32_t length) {
        LocalTensor<T> xrLocal = outQueueXR.DeQue<T>();
        DataCopy(xr_global[startpos], xrLocal, length); 
        outQueueXR.FreeTensor(xrLocal);

        LocalTensor<T> xwLocal = outQueueXW.DeQue<T>();
        DataCopy(xw_global[startpos], xwLocal, length); 
        outQueueXW.FreeTensor(xwLocal);

        LocalTensor<T> xkLocal = outQueueXK.DeQue<T>();
        DataCopy(xk_global[startpos], xkLocal, length); 
        outQueueXK.FreeTensor(xkLocal);

        LocalTensor<T> xvLocal = outQueueXV.DeQue<T>();
        DataCopy(xv_global[startpos], xvLocal, length); 
        outQueueXV.FreeTensor(xvLocal);
        PipeBarrier<PIPE_ALL>(); 

        LocalTensor<T> xaLocal = outQueueXA.DeQue<T>();
        DataCopy(xa_global[startpos], xaLocal, length); 
        outQueueXA.FreeTensor(xaLocal);
        PipeBarrier<PIPE_ALL>(); 

        LocalTensor<T> xgLocal = outQueueXG.DeQue<T>();
        DataCopy(xg_global[startpos], xgLocal, length); 
        outQueueXG.FreeTensor(xgLocal);
        PipeBarrier<PIPE_ALL>();
    }
    __aicore__ inline void CopyOutXTail(uint32_t startpos, uint32_t length) {
        LocalTensor<T> xrLocal = outQueueXR.DeQue<T>();
        LocalTensor<T> xwLocal = outQueueXW.DeQue<T>();
        LocalTensor<T> xkLocal = outQueueXK.DeQue<T>();
        LocalTensor<T> xvLocal = outQueueXV.DeQue<T>();
        LocalTensor<T> xaLocal = outQueueXA.DeQue<T>();
        LocalTensor<T> xgLocal = outQueueXG.DeQue<T>();
        for(int32_t k = 0; k < length; k++) {
            xr_global.SetValue(k + startpos, xrLocal.GetValue(k));
            xw_global.SetValue(k + startpos, xwLocal.GetValue(k));
            xk_global.SetValue(k + startpos, xkLocal.GetValue(k));
            xv_global.SetValue(k + startpos, xvLocal.GetValue(k));
            xa_global.SetValue(k + startpos, xaLocal.GetValue(k));
            xg_global.SetValue(k + startpos, xgLocal.GetValue(k));
        }
        PipeBarrier<PIPE_ALL>();
        outQueueXR.FreeTensor(xrLocal);
        outQueueXW.FreeTensor(xwLocal);
        outQueueXK.FreeTensor(xkLocal);
        outQueueXV.FreeTensor(xvLocal);
        outQueueXA.FreeTensor(xaLocal);
        outQueueXG.FreeTensor(xgLocal);

    }
    __aicore__ inline void CopyOutXUnalign(uint32_t startpos, uint32_t length) {
        uint32_t process_length = length / this->ALIGN_NUM * this->ALIGN_NUM;

        LocalTensor<T> xrLocal = outQueueXR.DeQue<T>();
        DataCopy(xr_global[startpos], xrLocal, process_length); 
        PipeBarrier<PIPE_ALL>();
        for(int32_t m = process_length; m < length; m++) {
            xr_global.SetValue(m + startpos, xrLocal.GetValue(m));
        }    
        PipeBarrier<PIPE_ALL>();
        outQueueXR.FreeTensor(xrLocal);

        LocalTensor<T> xwLocal = outQueueXW.DeQue<T>();
        DataCopy(xw_global[startpos], xwLocal, process_length); 
        PipeBarrier<PIPE_ALL>();
        for(int32_t m = process_length; m < length; m++) {
            xw_global.SetValue(m + startpos, xwLocal.GetValue(m));
        }     
        PipeBarrier<PIPE_ALL>();
        outQueueXW.FreeTensor(xwLocal);

        LocalTensor<T> xkLocal = outQueueXK.DeQue<T>();
        DataCopy(xk_global[startpos], xkLocal, process_length); 
        PipeBarrier<PIPE_ALL>();
        for(int32_t m = process_length; m < length; m++) {
            xk_global.SetValue(m + startpos, xkLocal.GetValue(m));
        }    
        PipeBarrier<PIPE_ALL>();
        outQueueXK.FreeTensor(xkLocal);

        LocalTensor<T> xvLocal = outQueueXV.DeQue<T>();
        DataCopy(xv_global[startpos], xvLocal, process_length); 
        PipeBarrier<PIPE_ALL>();
        for(int32_t m = process_length; m < length; m++) {
            xv_global.SetValue(m + startpos, xvLocal.GetValue(m));
        }      
        PipeBarrier<PIPE_ALL>();
        outQueueXV.FreeTensor(xvLocal);

        LocalTensor<T> xaLocal = outQueueXA.DeQue<T>();
        DataCopy(xa_global[startpos], xaLocal, process_length); 
        PipeBarrier<PIPE_ALL>();
        for(int32_t m = process_length; m < length; m++) {
            xa_global.SetValue(m + startpos, xaLocal.GetValue(m));
        }    
        PipeBarrier<PIPE_ALL>();
        outQueueXA.FreeTensor(xaLocal);

        LocalTensor<T> xgLocal = outQueueXG.DeQue<T>();
        DataCopy(xg_global[startpos], xgLocal, process_length); 
        PipeBarrier<PIPE_ALL>();
        for(int32_t m = process_length; m < length; m++) {
            xg_global.SetValue(m + startpos, xgLocal.GetValue(m));
        }     
        PipeBarrier<PIPE_ALL>();
        outQueueXG.FreeTensor(xgLocal);
    }
private:
    GlobalTensor<T> x_global;
    GlobalTensor<T> rwkvag_global;
    GlobalTensor<T> h0_global;
    GlobalTensor<T> xr_global, xw_global, xk_global, xv_global, xa_global, xg_global, ht_global; 

    TQueBind<TPosition::VECIN, TPosition::VECOUT, BUFFER_NUM> inQueueXBind;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueueRwkvag, inQueueH0;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueXNextLine;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueXR, outQueueXW, outQueueXK, outQueueXV, outQueueXA, outQueueXG;
    // TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueH0;
    TBuf<QuePosition::VECCALC> coor_range_Buf, voxel_size_Buf;
    TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmp2, QueueTmp3;
    TBuf<QuePosition::VECCALC> offset_buf;
    uint32_t B_length;
    uint32_t T_length;
    uint32_t C_length;
    uint32_t tiling_size;
    uint32_t core_rows;
    uint32_t core_rows_remain;
    uint32_t startPointer;
    uint32_t X_length_align;
    uint32_t ALIGN_NUM;

    LocalTensor<T> g_xLocal;
    AscendC::TPipe* pipe;
};

#include "token_shift_b1.h"
#include "token_shift_b1_align.h"
#include "token_shift_b1_t1_align.h"
#include "token_shift_b1_t1_align_1core.h"

extern "C" __global__ __aicore__ void token_shift(GM_ADDR x, GM_ADDR rwkvag, GM_ADDR h0, GM_ADDR xr, GM_ADDR xw, GM_ADDR xk, GM_ADDR xv, GM_ADDR xa, GM_ADDR xg, GM_ADDR ht, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    
    if (TILING_KEY_IS(100)) {
        TPipe pipe;
        TokenShift<DTYPE_X> op;
        op.Init(x, rwkvag, h0, 
                xr, xw, xk, xv, xa, xg, ht,
                workspace,
                tiling_data.B_length,
                tiling_data.T_length,
                tiling_data.C_length,
                tiling_data.tiling_size,
                tiling_data.core_rows,
                tiling_data.core_rows_remain,
                tiling_data.X_length_align,
                &pipe);
        op.Process();
        // op.ProcessHT(x);
    } else if (TILING_KEY_IS(200)) {
        TPipe pipe;
        TokenShiftB1<DTYPE_X> op;
        op.Init(x, rwkvag, h0, 
                xr, xw, xk, xv, xa, xg, ht,
                workspace,
                tiling_data.aivNum,
                tiling_data.B_length,
                tiling_data.T_length,
                tiling_data.C_length,
                tiling_data.tiling_size,
                tiling_data.core_rows,
                tiling_data.core_rows_remain,
                tiling_data.X_length_align,
                &pipe);
        op.Process();
    } else if (TILING_KEY_IS(201)) {
        TPipe pipe;
        TokenShiftB1Align<DTYPE_X> op;
        op.Init(x, rwkvag, h0, 
                xr, xw, xk, xv, xa, xg, ht,
                workspace,
                tiling_data.aivNum,
                tiling_data.B_length,
                tiling_data.T_length,
                tiling_data.C_length,
                tiling_data.tiling_size,
                tiling_data.core_rows,
                tiling_data.core_rows_remain,
                tiling_data.X_length_align,
                &pipe);
        op.Process();
    } else if (TILING_KEY_IS(202)) {
        TPipe pipe;
        TokenShiftB1T1Align_1Core<DTYPE_X> op;
        op.Init(x, rwkvag, h0, 
                xr, xw, xk, xv, xa, xg, ht,
                workspace,
                tiling_data.aivNum,
                tiling_data.B_length,
                tiling_data.T_length,
                tiling_data.C_length,
                tiling_data.tiling_size,
                tiling_data.core_rows,
                tiling_data.core_rows_remain,
                tiling_data.X_length_align,
                &pipe);
        op.Process();
    } else if (TILING_KEY_IS(204)) {
        TokenShiftB1T1Align<DTYPE_X> op;
        op.Init(x, rwkvag, h0, 
                xr, xw, xk, xv, xa, xg, ht,
                workspace,
                tiling_data.aivNum,
                tiling_data.B_length,
                tiling_data.T_length,
                tiling_data.C_length,
                tiling_data.tiling_size,
                tiling_data.core_rows,
                tiling_data.core_rows_remain,
                tiling_data.X_length_align);
        op.Process();
    }
}